import glob
import os
import shutil

import cv2
import numpy as np
import torch
from PIL import Image
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms

from data_loader import RescaleT
from data_loader import SalObjDataset
from data_loader import ToTensorLab

# 尝试导入U2Net（如果可用）
try:
    from model import U2NET, U2NETP

    U2NET_AVAILABLE = True
except ImportError:
    U2NET_AVAILABLE = False
    print("警告: 未找到U2Net库，U2Net功能将不可用")

class U2netImageMatting:
    def __init__(self, input_path, output_path, use_gpu=True, model_name='u2netp'):
        self.input_path = input_path
        self.output_path = output_path
        self.save_path = output_path
        self.use_gpu = use_gpu and torch.cuda.is_available()
        self.model_name = model_name
        self.model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth')

        if os.path.isfile(self.input_path):
            self.img_name_list = [self.input_path]
        elif os.path.isdir(self.input_path):
            # 获取所有文件（不含文件夹）
            self.img_name_list = [
                os.path.join(self.input_path, name)
                for name in os.listdir(self.input_path)
                if os.path.isfile(os.path.join(self.input_path, name))
            ]
        else:
            raise ValueError("Input path is neither a valid file nor a valid directory.")

    # normalize the predicted SOD probability map
    def normPRED(self, d):
        ma = torch.max(d)
        mi = torch.min(d)

        dn = (d - mi) / (ma - mi)

        return dn

    def save_output(self, image_name, pred, d_dir):
        predict = pred
        predict = predict.squeeze()
        predict_np = predict.cpu().data.numpy()

        # 读取原始图像
        original_image = Image.open(image_name).convert("RGBA")
        width, height = original_image.size

        # 调整预测掩码的大小以匹配原始图像
        mask = Image.fromarray(predict_np * 255).convert("L")
        mask = mask.resize((width, height), resample=Image.BILINEAR)

        # 将掩码应用于原始图像
        result = Image.new("RGBA", (width, height))
        for x in range(width):
            for y in range(height):
                r, g, b, a = original_image.getpixel((x, y))
                alpha = mask.getpixel((x, y))
                result.putpixel((x, y), (r, g, b, alpha))

        imge_name = image_name.split(os.sep)[-1]
        result_name = imge_name.split('/')[-1]
        result_name, result_type = result_name.split('.')
        self.save_path = d_dir + f'/{result_name}_matting_{self.model_name}.png'
        result.save(self.save_path)

    def process_images(self):
        # --------- 2. dataloader ---------
        # 1. dataloader
        test_salobj_dataset = SalObjDataset(img_name_list=self.img_name_list,
                                            lbl_name_list=[],
                                            transform=transforms.Compose([RescaleT(320),
                                                                          ToTensorLab(flag=0)])
                                            )
        test_salobj_dataloader = DataLoader(test_salobj_dataset,
                                            batch_size=1,
                                            shuffle=False,
                                            num_workers=1)

        # --------- 3. model define ---------
        if self.model_name == 'u2net':
            # print("...load U2NET---173.6 MB")
            net = U2NET(3, 1)
        elif self.model_name == 'u2netp':
            # print("...load U2NEP---4.7 MB")
            net = U2NETP(3, 1)

        if self.use_gpu:
            net.load_state_dict(torch.load(self.model_dir, weights_only=True))
            net.cuda()
        else:
            net.load_state_dict(torch.load(self.model_dir, map_location='cpu', weights_only=True))
        net.eval()

        # --------- 4. inference for each image ---------
        for i_test, data_test in enumerate(test_salobj_dataloader):

            # print("inferencing:", self.img_name_list[i_test].split(os.sep)[-1])

            inputs_test = data_test['image']
            inputs_test = inputs_test.type(torch.FloatTensor)

            if self.use_gpu:
                inputs_test = Variable(inputs_test.cuda())
            else:
                inputs_test = Variable(inputs_test)

            d1, d2, d3, d4, d5, d6, d7 = net(inputs_test)

            # normalization
            pred = d1[:, 0, :, :]
            pred = self.normPRED(pred)

            # save results to test_results folder
            if not os.path.exists(self.output_path):
                os.makedirs(self.output_path, exist_ok=True)
            self.save_output(self.img_name_list[i_test], pred, self.output_path)

            del d1, d2, d3, d4, d5, d6, d7

class MattingProcessor:
    """图像抠图处理类，提供图像抠图功能"""


    def __init__(self):
        # 模型缓存
        self.models = {}
        self.u2netImageMatting = None

    def process_image_u2net(self, input_path, output_dir, model_name, use_gpu=False):
        """使用U2Net处理单个图像并保存结果

        Args:
            input_path: 输入图像路径
            output_dir: 输出目录
            model_name: 模型名称 ('u2net' 或 'u2netp')
            use_gpu: 是否使用GPU

        Returns:
            处理后的图像保存路径
        """
        self.u2netImageMatting = U2netImageMatting(input_path, output_dir, use_gpu, model_name)
        self.u2netImageMatting.process_images()
        return self.u2netImageMatting.save_path

    def process_image_grabcut(self, input_path, output_dir):
        """使用OpenCV GrabCut处理单个图像并保存结果

        Args:
            input_path: 输入图像路径
            output_dir: 输出目录

        Returns:
            处理后的图像保存路径
        """
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)


        try:
            # 优先使用Pillow读取图像
            pil_img = Image.open(input_path)
            image = np.array(pil_img)

            # 转换为OpenCV的BGR格式
            if len(image.shape) == 3:
                if image.shape[2] == 3:
                    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
                elif image.shape[2] == 4:
                    image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGR)
            else:
                # 灰度图像
                image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
            # # 读取图像
            # image = cv2.imread(input_path)
            # if image is None:
            #     raise ValueError(f"无法读取图像: {input_path}")

            # 转换为RGB
            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

            # 创建掩码和模型
            mask = np.zeros(image_rgb.shape[:2], np.uint8)
            bgd_model = np.zeros((1, 65), np.float64)
            fgd_model = np.zeros((1, 65), np.float64)

            # 定义矩形区域（自动计算）
            height, width = image_rgb.shape[:2]
            rect = (50, 50, width - 100, height - 100)

            # 应用GrabCut算法
            cv2.grabCut(image_rgb, mask, rect, bgd_model, fgd_model, 5, cv2.GC_INIT_WITH_RECT)

            # 生成透明度掩码
            mask2 = np.where((mask == 2) | (mask == 0), 0, 255).astype('uint8')

            # 确定输出格式
            filename, ext = os.path.splitext(os.path.basename(input_path))
            output_ext = '.png' if ext.lower() in ['.jpg', '.jpeg', '.png'] else ext
            output_filename = f"{filename}_matting_grabcut{output_ext}"
            output_path = os.path.join(output_dir, output_filename)

            # 根据格式保存图像
            if output_ext.lower() == '.png':
                # 保存带透明通道的PNG
                result = np.zeros((height, width, 4), dtype=np.uint8)
                result[:, :, :3] = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                result[:, :, 3] = mask2
                pil_img = Image.fromarray(result)
                pil_img.save(output_path)
            else:
                # 保存不透明图像（如JPG）
                result = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                # 应用前景掩码
                result = cv2.bitwise_and(result, result, mask=mask2)
                # 转换回BGR用于OpenCV保存
                result_bgr = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
                cv2.imwrite(output_path, result_bgr)

            return output_path

        except Exception as e:
            print(f"处理图像时出错 ({input_path}): {str(e)}")
            # 复制原始图像作为备份
            backup_path = os.path.join(output_dir, f"{filename}_backup{ext}")
            shutil.copy2(input_path, backup_path)
            raise